"""
BaseAgent 类 - 交易代理的基类
封装了核心功能，包括 MCP 工具管理、AI 代理创建和交易执行
"""

import os
import json
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
import traceback, asyncio

# 导入必要的库和模块
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langchain.agents import create_agent
from dotenv import load_dotenv

# 导入项目工具
import sys

from symbols import all_symbols
from tools.trade_date_tools import TradingCalendar

# 获取项目根目录并添加到 Python 路径中
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

# 导入工具函数和常量
from tools.general_tools import extract_conversation, extract_tool_messages, get_config_value, write_config_value
from tools.price_tools import add_no_trade_record
from prompts.agent_prompt import get_agent_system_prompt, STOP_SIGNAL

# 加载环境变量
load_dotenv()


class BaseAgent:
    """
    交易代理的基类
    
    主要功能：
    1. MCP 工具管理和连接
    2. AI 代理创建和配置
    3. 交易执行和决策循环
    4. 日志记录和管理
    5. 持仓和配置管理
    """

    def __init__(
        self,
        signature: str,
        basemodel: str,
        stock_symbols: Optional[List[str]] = None,
        mcp_config: Optional[Dict[str, Dict[str, Any]]] = None,
        log_path: Optional[str] = None,
        max_steps: int = 10,
        max_retries: int = 3,
        base_delay: float = 0.5,
        openai_base_url: Optional[str] = None,
        initial_cash: float = 10000.0,
        init_date: str = "2025-10-13"
    ):
        """
        初始化 BaseAgent
        
        参数:
            signature: 代理的签名/名称
            basemodel: 基础模型名称
            stock_symbols: 股票代码列表，默认为 NASDAQ 100
            mcp_config: MCP 工具配置，包含端口和 URL 信息
            log_path: 日志路径，默认为 ./data/agent_data
            max_steps: 最大推理步数
            max_retries: 最大重试次数
            base_delay: 重试基础延迟时间
            openai_base_url: OpenAI API 基础 URL
            initial_cash: 初始现金金额
            init_date: 初始化日期
        """
        # 设置基本属性
        self.signature = signature
        self.basemodel = basemodel
        self.stock_symbols = stock_symbols or all_symbols
        self.max_steps = max_steps
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.initial_cash = initial_cash
        self.init_date = init_date
        
        # 设置 MCP 配置
        self.mcp_config = mcp_config or self._get_default_mcp_config()
        
        # 设置日志路径
        self.base_log_path = log_path or "./data/agent_data"
        
        # 设置 OpenAI 配置
        self.openai_base_url = openai_base_url or os.getenv("OPENAI_API_BASE")

        # 根据不同的模型名称，读取api_key
        if "deepseek" in self.basemodel.lower():
            self.api_key = os.getenv("DEEPSEEK_API_KEY")
        elif "qwen" in self.basemodel.lower():
            self.api_key = os.getenv("DASHSCOPE_API_KEY")
        else:
            self.api_key = os.getenv("OPENAI_API_KEY")
        
        # 初始化组件
        self.client: Optional[MultiServerMCPClient] = None
        self.tools: Optional[List] = None
        self.model: Optional[ChatOpenAI] = None
        self.agent: Optional[Any] = None
        
        # 数据路径
        self.data_path = os.path.join(self.base_log_path, self.signature)
        self.position_file = os.path.join(self.data_path, "position", "position.jsonl")
        
    def _get_default_mcp_config(self) -> Dict[str, Dict[str, Any]]:
        """获取默认的 MCP 配置"""
        return {
            "math": {
                "transport": "streamable_http",
                "url": f"http://localhost:{os.getenv('MATH_HTTP_PORT', '8000')}/mcp",
            },
            "stock_local": {
                "transport": "streamable_http",
                "url": f"http://localhost:{os.getenv('GETPRICE_HTTP_PORT', '8003')}/mcp",
            },
            "search": {
                "transport": "streamable_http",
                "url": f"http://localhost:{os.getenv('SEARCH_HTTP_PORT', '8001')}/mcp",
            },
            "trade": {
                "transport": "streamable_http",
                "url": f"http://localhost:{os.getenv('TRADE_HTTP_PORT', '8002')}/mcp",
            },
        }
    
    async def initialize(self) -> None:
        """初始化 MCP 客户端和 AI 模型"""
        print(f"🚀 初始化代理: {self.signature}")
        
        # 创建 MCP 客户端
        self.client = MultiServerMCPClient(self.mcp_config)
        
        # 获取工具
        try:
            self.tools = await self.client.get_tools()
        except* Exception as eg:  # Python 3.11+ ExceptionGroup
            for e in eg.exceptions:
                traceback.print_exception(type(e), e, e.__traceback__)
            raise eg
        print(f"✅ 已加载 {len(self.tools)} 个 MCP 工具")
        
        # 创建 AI 模型
        self.model = ChatOpenAI(
            model=self.basemodel,
            base_url=self.openai_base_url,
            api_key=self.api_key,
            max_retries=3,
            timeout=30
        )
        
        # 注意：代理将在 run_trading_session() 中基于特定日期创建
        # 因为系统提示需要当前日期和价格信息
        
        print(f"✅ 代理 {self.signature} 初始化完成")
    
    def _setup_logging(self, today_date: str) -> str:
        """设置日志文件路径"""
        # 创建基于日期的日志路径
        log_path = os.path.join(self.base_log_path, self.signature, 'log', today_date)
        # 如果目录不存在则创建
        if not os.path.exists(log_path):
            os.makedirs(log_path)
        # 返回日志文件的完整路径
        return os.path.join(log_path, "log.jsonl")
    
    def _log_message(self, log_file: str, new_messages: List[Dict[str, str]]) -> None:
        """将消息记录到日志文件中"""
        # 创建日志条目
        log_entry = {
            "timestamp": datetime.now().isoformat(),  # 当前时间戳
            "signature": self.signature,  # 代理签名
            "new_messages": new_messages  # 要记录的消息列表
        }
        # 追加写入日志文件
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(log_entry, ensure_ascii=False) + "\n")
    
    async def _ainvoke_with_retry(self, message: List[Dict[str, str]]) -> Any:
        """带重试机制的代理调用"""
        # 尝试多次调用代理
        for attempt in range(1, self.max_retries + 1):
            try:
                # 调用代理并返回结果
                return await self.agent.ainvoke(
                    {"messages": message}, 
                    {"recursion_limit": 100}  # 设置递归限制
                )
            except Exception as e:
                # 如果是最后一次尝试，则抛出异常
                if attempt == self.max_retries:
                    raise e
                # 否则等待一段时间后重试
                print(f"⚠️ 尝试 {attempt} 失败，{self.base_delay * attempt} 秒后重试...")
                print(f"错误详情: {e}")
                await asyncio.sleep(self.base_delay * attempt)
    
    async def run_trading_session(self, today_date: str) -> None:
        """
        运行单日交易会话
        
        参数:
            today_date: 交易日期
        """
        #print(f"📈 开始交易会话: {today_date}")
        
        # 设置日志
        log_file = self._setup_logging(today_date)
        
        # 更新系统提示并创建代理
        self.agent = create_agent(
            self.model,
            tools=self.tools,
            system_prompt=get_agent_system_prompt(today_date, self.signature),
        )
        
        # 初始用户查询
        user_query = [{"role": "user", "content": f"请分析并更新今天({today_date})的持仓情况。"}]
        message = user_query.copy()
        
        # 记录初始消息
        self._log_message(log_file, user_query)
        
        # 交易循环
        current_step = 0
        while current_step < self.max_steps:
            current_step += 1
            print(f"🔄 步骤 {current_step}/{self.max_steps}")
            
            try:
                # 调用代理
                response = await self._ainvoke_with_retry(message)
                
                # 提取代理响应
                agent_response = extract_conversation(response, "final")
                
                # 检查停止信号
                if STOP_SIGNAL in agent_response:
                    print("✅ 收到停止信号，交易会话结束")
                    print(agent_response)
                    self._log_message(log_file, [{"role": "assistant", "content": agent_response}])
                    break
                
                # 提取工具消息
                tool_msgs = extract_tool_messages(response)
                tool_response = '\n'.join([msg.content for msg in tool_msgs])
                
                # 准备新消息
                new_messages = [
                    {"role": "assistant", "content": agent_response},
                    {"role": "user", "content": f'工具结果: {tool_response}'}
                ]
                
                # 添加新消息到对话历史
                message.extend(new_messages)
                
                # 记录消息
                self._log_message(log_file, new_messages[0])
                self._log_message(log_file, new_messages[1])
                
            except Exception as e:
                print(f"❌ 交易会话错误: {str(e)}")
                print(f"错误详情: {e}")
                raise
        
        # 处理交易结果
        await self._handle_trading_result(today_date)
    
    async def _handle_trading_result(self, today_date: str) -> None:
        """处理交易结果"""
        # 获取交易标志
        if_trade = get_config_value("IF_TRADE")
        
        if if_trade:
            # 如果执行了交易，重置标志并记录
            write_config_value("IF_TRADE", False)
            print("✅ 交易完成")
        else:
            # 如果没有交易，添加无交易记录
            print("📊 无交易，维持现有持仓")
            try:
                add_no_trade_record(today_date, self.signature)
            except NameError as e:
                print(f"❌ 名称错误: {e}")
                raise
            write_config_value("IF_TRADE", False)
    
    def register_agent(self) -> None:
        """注册新代理，创建初始持仓"""
        # 检查持仓文件是否已存在
        if os.path.exists(self.position_file):
            print(f"⚠️ 持仓文件 {self.position_file} 已存在，跳过注册")
            return
        
        # 确保目录结构存在
        position_dir = os.path.join(self.data_path, "position")
        if not os.path.exists(position_dir):
            os.makedirs(position_dir)
            print(f"📁 创建持仓目录: {position_dir}")
        
        # 创建初始持仓
        init_position = {symbol: 0 for symbol in self.stock_symbols}  # 所有股票初始持仓为0
        init_position['CASH'] = self.initial_cash  # 设置初始现金
        
        # 写入初始持仓信息到文件
        with open(self.position_file, "w") as f:  # 使用"w"模式确保创建新文件
            f.write(json.dumps({
                "date": self.init_date, 
                "id": 0, 
                "positions": init_position
            }) + "\n")
        
        # 打印注册完成信息
        print(f"✅ 代理 {self.signature} 注册完成")
        print(f"📁 持仓文件: {self.position_file}")
        print(f"💰 初始现金: ${self.initial_cash}")
        print(f"📊 股票数量: {len(self.stock_symbols)}")
    
    def get_trading_dates(self, init_date: str, end_date: str) -> List[str]:
        """
        根据已有持仓记录，计算从“上次已处理日期的次日”到 end_date 之间的所有工作日，
        返回需要补跑的交易日列表；若持仓文件不存在则先注册代理并把 init_date 视为起点。
        
        参数:
            init_date: 开始日期
            end_date: 结束日期
            
        返回:
            交易日列表
        """
        max_date = None
        
        # 如果持仓文件不存在，则注册代理并设置最大日期为初始日期
        if not os.path.exists(self.position_file):
            self.register_agent()
            max_date = init_date
        else:
            # 读取现有持仓文件，找到最新日期
            with open(self.position_file, "r") as f:
                for line in f:
                    doc = json.loads(line)
                    current_date = doc['date']
                    if max_date is None:
                        max_date = current_date
                    else:
                        # 比较日期，保留较大的日期
                        current_date_obj = datetime.strptime(current_date, "%Y-%m-%d")
                        max_date_obj = datetime.strptime(max_date, "%Y-%m-%d")
                        if current_date_obj > max_date_obj:
                            max_date = current_date
        if max_date is None:
            max_date = init_date
        # 检查是否需要处理新日期
        max_date_obj = datetime.strptime(max_date, "%Y-%m-%d")
        end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
        
        if end_date_obj <= max_date_obj:
            return []  # 没有需要处理的新日期
        
        # 生成交易日列表
        trading_dates = []
        current_date = max_date_obj + timedelta(days=1)
        # 创建交易日历实例（会自动缓存）
        calendar = TradingCalendar()

        while current_date <= end_date_obj:
            #if current_date.weekday() < 5:  # 只包括工作日
            #    trading_dates.append(current_date.strftime("%Y-%m-%d"))
            result = calendar.is_trading_day(current_date.strftime("%Y-%m-%d"))
            if result:
                trading_dates.append(current_date.strftime("%Y-%m-%d"))

            current_date += timedelta(days=1)
        
        return trading_dates
    
    async def run_with_retry(self, today_date: str) -> None:
        """带重试机制的运行方法"""
        # 尝试多次运行交易会话
        for attempt in range(1, self.max_retries + 1):
            try:
                print(f"🔄 尝试运行 {self.signature} - {today_date} (尝试 {attempt})")
                await self.run_trading_session(today_date)
                print(f"✅ {self.signature} - {today_date} 运行成功")
                return
            except Exception as e:
                print(f"❌ 尝试 {attempt} 失败: {str(e)}")
                if attempt == self.max_retries:
                    print(f"💥 {self.signature} - {today_date} 所有重试均失败")
                    raise
                else:
                    # 计算等待时间并进行下一次尝试
                    wait_time = self.base_delay * attempt
                    print(f"⏳ {wait_time} 秒后重试...")
                    await asyncio.sleep(wait_time)
    
    async def run_date_range(self, init_date: str, end_date: str) -> None:
        """
        运行日期范围内的所有交易日，包括从“上次已处理日期的次日”到 end_date 之间的所有工作日。
        
        参数:
            init_date: 开始日期
            end_date: 结束日期
        """
        print(f"📅 运行日期范围: {init_date} 到 {end_date}")
        
        # 获取交易日列表
        trading_dates = self.get_trading_dates(init_date, end_date)
        
        if not trading_dates:
            print(f"ℹ️ 没有需要处理的交易日")
            return
        
        print(f"📊 需要处理的交易日: {trading_dates}")
        
        # 处理每个交易日
        for date in trading_dates:
            #print(f"🔄 处理 {self.signature} - 日期: {date}")
            
            # 设置配置
            write_config_value("TODAY_DATE", date)
            write_config_value("SIGNATURE", self.signature)
            
            try:
                await self.run_with_retry(date)
            except Exception as e:
                print(f"❌ 处理 {self.signature} - 日期: {date} 时出错")
                print(e)
                raise
        
        print(f"✅ {self.signature} 处理完成")
    
    def get_position_summary(self) -> Dict[str, Any]:
        """获取持仓摘要"""
        # 检查持仓文件是否存在
        if not os.path.exists(self.position_file):
            return {"error": "持仓文件不存在"}
        
        # 读取所有持仓记录
        positions = []
        with open(self.position_file, "r") as f:
            for line in f:
                positions.append(json.loads(line))
        
        # 检查是否有持仓记录
        if not positions:
            return {"error": "没有持仓记录"}
        
        # 获取最新的持仓记录
        latest_position = positions[-1]
        # 返回摘要信息
        return {
            "signature": self.signature,
            "latest_date": latest_position.get("date"),
            "positions": latest_position.get("positions", {}),
            "total_records": len(positions)
        }
    
    def __str__(self) -> str:
        """返回对象的字符串表示"""
        return f"BaseAgent(signature='{self.signature}', basemodel='{self.basemodel}', stocks={len(self.stock_symbols)})"
    
    def __repr__(self) -> str:
        """返回对象的正式字符串表示"""
        return self.__str__()
